"""
Create: JiBingyu 2024-08-10
Description: 将比赛的USV环境模型进行封装，方便调用，使用Gym结构
"""

import ctypes
import os
import time
import unittest
from typing import List, Optional
from unittest.mock import patch

import gymnasium as gym
import matplotlib.pyplot as plt
import numpy as np
from gymnasium import spaces

from Function_j import *
from Function_j import suppress_stdout_stderr
from MultiUSVModel import MultiUSVModel
from StructType import *


class USV_Env(gym.Env):
    def __init__(
        self,
        num_usv,
        render_mode: Optional[str] = None,
    ):
        super(USV_Env, self).__init__()
        """
        初始化函数
        :param usv_num: USV数量
        :param map_path: 地图文件路径
        :param config_path: 配置文件路径
        :param show_flag: 是否显示地图
        """
        # 定义动作空间：速度和角度变化
        self.action_space = spaces.Box(
            low=np.array([3, -np.pi / 18] * num_usv),  # 最小速度变化和最小角度变化
            high=np.array([15, np.pi / 18] * num_usv),  # 最大速度变化和最大角度变化
            dtype=np.float32,
        )

        # 定义状态空间：位置(x, y)和方向角度
        self.observation_space = spaces.Box(
            low=np.array(
                [0.0, 0.0, 0.0, -np.pi, 1] * 2 * num_usv
            ),  # 最小位置和最小角度
            high=np.array(
                [100.0, 100.0, 1.0, np.pi, 2] * 2 * num_usv
            ),  # 最大位置和最大角度
            dtype=np.float32,
        )

        self.render_mode = render_mode

        self.num_usv = num_usv
        # 初始化USV模型
        os.add_dll_directory(os.path.join(os.getcwd(), "./dll"))
        # 初始化
        self.cmultilUSV = MultiUSVModel(
            dll_path=r"./dll/usvmodel.dll",
            dll_deps=[os.path.join(r"./dll", dll_i) for dll_i in os.listdir(r"./dll")],
        )

        # 态势信息
        # 控制模式，1 - 航速航向；2 - 航路，建议使用航速航向
        self.flagModel = 1
        self.curOBSFusion_data = ctypes.POINTER(
            InputOBSStruct
        )()  # 指向InputOBSStruct类型的指针
        self.curOBSFusion_size = ctypes.c_int()  # 障碍物数量
        # 导航信息
        self.multiUSVStatus_data = ctypes.POINTER(USVOutPutStructType)()
        self.multiUSVStatus_size = ctypes.c_int()

        self.isEnd_int = ctypes.c_int()
        # 任务航路
        self.taskPath_data = ctypes.POINTER(SFormationTaskInfo)()
        self.taskPath_size = ctypes.c_int()
        # 构型信息
        self.referInfo_data = ctypes.POINTER(referUSVInfo)()
        self.referInfo_size = ctypes.c_int()

        ################################################################################
        self.Task_data = []
        self.Refe_data = []
        self.Obs_data = []
        self.USV_data = []

        self.reset()

    @patch("builtins.open", side_effect=Exception("File operation is not allowed"))
    def reset(self, seed=0):
        """
        重置环境
        """

        self.done = [False] * self.num_usv
        self.usv_step = np.zeros((self.num_usv))

        # 各艇控制指令
        self.multiUSVInitInfo: List[USVInitStructType] = []
        # 整个任务迭代次数
        self.iterNum = 0

        # 算法更新后清零
        self.algIterNum = 0

        with suppress_stdout_stderr():
            self.cmultilUSV.multiInitialize(self.flagModel)
            self.cmultilUSV.updateMultiUSVControl(
                (USVInitStructType * len(self.multiUSVInitInfo))(
                    *self.multiUSVInitInfo
                ),
                len(self.multiUSVInitInfo),
                self.iterNum,
                self.flagModel,
            )
            # 获取导航信息和态势信息
            self.cmultilUSV.getMultiNaviObsInfo(
                self.iterNum,
                ctypes.byref(self.multiUSVStatus_data),
                ctypes.byref(self.multiUSVStatus_size),
                ctypes.byref(self.curOBSFusion_data),
                ctypes.byref(self.curOBSFusion_size),
            )
            # 获取任务信息
            rcvFlag = self.cmultilUSV.getTaskInfo(
                ctypes.byref(self.isEnd_int),
                ctypes.byref(self.taskPath_data),
                ctypes.byref(self.taskPath_size),
                ctypes.byref(self.referInfo_data),
                ctypes.byref(self.referInfo_size),
                self.iterNum,
            )
        _, self.state = self.translate_information()
        self.state = self.state.reshape(-1)
        self.next_state = None
        return self.state, {}

    def step(self, actions):
        """
        给定飞机的速度和角度，根据平台的动力学模型，计算下一个状态
        """
        ##########################################更新控制指令###########################
        # 更新控制指令
        done = False
        with suppress_stdout_stderr():
            self.cmultilUSV.updateMultiUSVControl(
                (USVInitStructType * len(self.multiUSVInitInfo))(
                    *self.multiUSVInitInfo
                ),
                len(self.multiUSVInitInfo),
                self.iterNum,
                self.flagModel,
            )
            # 获取导航信息和态势信息
            self.cmultilUSV.getMultiNaviObsInfo(
                self.iterNum,
                ctypes.byref(self.multiUSVStatus_data),
                ctypes.byref(self.multiUSVStatus_size),
                ctypes.byref(self.curOBSFusion_data),
                ctypes.byref(self.curOBSFusion_size),
            )
            # 获取任务信息
            rcvFlag = self.cmultilUSV.getTaskInfo(
                ctypes.byref(self.isEnd_int),
                ctypes.byref(self.taskPath_data),
                ctypes.byref(self.taskPath_size),
                ctypes.byref(self.referInfo_data),
                ctypes.byref(self.referInfo_size),
                self.iterNum,
            )
        # 触发任务
        # if 1 == rcvFlag:
        #     if 1 == self.isEnd_int.value:  # 到达最后一个任务点，不再发送任务
        #         self.cmultilUSV.taskEnd(self.iterNum, self.flagModel)
        #         done = True  # 跳出算法循环
        #     algIterNum = 0  # 新任务，次数清零
        ##############################################End###############################

        # 算法更新控制指令
        self.translate_action(actions)
        # 计算奖励
        reward = self.caculate_reward()
        # 更换状态

        # 在self.done中，False的索引位置self.num_usv加1
        for i in range(self.num_usv):
            if not self.done[i]:
                self.usv_step[i] += 1
        self.state = self.next_state.reshape(-1)

        if self.render_mode == "human":
            self.render()
        return self.state, reward, done, False, {}

    def render(self, mode="human"):
        if not hasattr(self, "fig"):
            self.fig, self.ax = plt.subplots()
            self.ax.set_xlim(0, 1000)
            self.ax.set_ylim(-1000, 1000)
            (self.usv,) = self.ax.plot([], [], "bo", markersize=10)
            (self.target,) = self.ax.plot([], [], "*", markersize=10)
            # 输出速度和角度
            self.speed_text = self.ax.text(0, 0, "")
            self.angle_text = self.ax.text(0, 30, "")

        self.usv.set_data(
            [x for x, _, _, _, _ in self.USV_data],
            [y for _, y, _, _, _ in self.USV_data],
        )
        self.target.set_data(
            [300 + x for x, _, _, _, _ in self.Refe_data],
            [300 + y for _, y, _, _, _ in self.Refe_data],
        )
        # self.speed_text.set_text(
        #     f"speed:{self.policy[0]:.2f}, clip: {self.state[2]:.2f}"
        # )
        # self.angle_text.set_text(
        #     f"angle: {self.policy[1]:.2f}, clip:{self.state[3]:.2f}"
        # )
        plt.draw()
        plt.pause(0.01)

    def translate_action(self, actions):
        """
        输入：action
            [
                [速度, 角度],
                [速度, 角度],
                ...]
            ]
        输出：控制指令
        将动作转换为控制指令
        """
        multiUSVStatus_list, self.next_state = self.translate_information()
        # 根据实时位置，任务信息求下一时刻艇的期望航速航向
        self.multiUSVInitInfo: List[USVInitStructType] = []
        for i, it in enumerate(multiUSVStatus_list):
            updateUSVControl = USVInitStructType()
            updateUSVControl.iUSVId = it.iUSVId
            updateUSVControl.dbInitX = it.dbCurX
            updateUSVControl.dbInitY = it.dbCurY

            #################################### 控制算法部分 ###############################
            updateUSVControl.fInitVel, updateUSVControl.fInitYaw = (
                actions[2 * i],
                90 - actions[1 + 2 * i],
            )
            self.multiUSVInitInfo.append(updateUSVControl)
        return 0

    def translate_information(self):
        """
        输入：无
        输出：将C++中的数据转换为Python中的数据
        计算数据
        """
        taskPath_list: List[SFormationTaskInfo] = [
            self.taskPath_data[index] for index in range(self.taskPath_size.value)
        ]
        referInfo_list: List[referUSVInfo] = [
            self.referInfo_data[index] for index in range(self.referInfo_size.value)
        ]
        curOBSFusion_list: List[InputOBSStruct] = [
            self.curOBSFusion_data[index]
            for index in range(self.curOBSFusion_size.value)
        ]
        multiUSVStatus_list: List[USVOutPutStructType] = [
            self.multiUSVStatus_data[index]
            for index in range(self.multiUSVStatus_size.value)
        ]
        if len(taskPath_list) != 0:
            self.Task_data = []
            for it in taskPath_list:
                self.Task_data.append([it.referX, it.referY])
        if len(referInfo_list) != 0:
            self.Refe_data = []
            for it in referInfo_list:
                self.Refe_data.append(
                    [
                        it.relatDis * np.sin(np.radians(90 - it.relatYaw)),
                        it.relatDis * np.cos(np.radians(90 - it.relatYaw)),
                        it.relatYaw,
                        it.type,
                        it.USVID,
                    ]
                )
        if len(curOBSFusion_list) != 0:
            self.Obs_data = []
            for it in curOBSFusion_list:
                # 这里的 obsR 和 fOBSV与说明书上反的
                self.Obs_data.append([it.dOBSX, it.dOBSY, it.fOBSV])
        if len(multiUSVStatus_list) != 0:
            self.USV_data = []
            for it in multiUSVStatus_list:
                self.USV_data.append(
                    [it.dbCurX, it.dbCurY, it.fCurVel, it.fCurYaw, it.iUSVId]
                )
        # print("ReferInfo", len(referInfo_list))

        Information = np.concatenate([self.Refe_data + self.USV_data], axis=0)
        return multiUSVStatus_list, Information

    def caculate_reward(self):
        """
        计算奖励
        """
        # time.sleep(0.001)
        # print(self.state.shape)
        # print(self.next_state.shape)
        # print(self.Refe_data)
        if self.Refe_data[0][-2] == 0:
            # 没有固定分配位置
            reward = self.reward_no_task()
        else:
            # 有固定分配位置
            reward = self.reward_task()
        return reward

    def reward_no_task(self):
        """
        没有固定分配位置
        """
        pass

    def reward_task(self):
        """
        固定分配位置
        """
        Information = np.copy(self.next_state)
        ref_info = None
        reward = 0
        for i in range(self.num_usv):
            if Information[i][-2] == 1:
                ref_info = Information[i]
        for i in range(self.num_usv):
            Information[i][0] += ref_info[0]
            Information[i][1] += ref_info[1]

        sorted_indices = np.argsort(Information[:, -1])
        Information = Information[sorted_indices]
        for i in range(0, self.num_usv * 2, 2):
            distance = np.sqrt(
                (Information[i][0] - Information[i + 1][0]) ** 2
                + (Information[i][1] - Information[i + 1][1]) ** 2
            )
            if distance < 50:
                if not self.done[i / 2]:
                    self.done[i / 2] = True
                    reward += 100 / self.usv_step[i / 2]
            else:
                reward += -distance
        return reward


if __name__ == "__main__":
    env = USV_Env(6)
    for _ in range(30):
        # print("reset—")
        s, _ = env.reset()
        # print(s.shape)
        # print("reset")
        for i in range(100):
            action = env.action_space.sample()
            env.step([3, 0] * 6)
        print(_, i)
    env.close()
    del env
